import time

import mujoco
import mujoco.viewer
import numpy as np
from etils import epath

from mujoco_playground._src.locomotion.g1.base import get_assets

model = mujoco.MjModel.from_xml_string(epath.Path('/home/huang/mujoco_playground/mujoco_playground/_src/locomotion/g1/xmls/scene_mjx_feetonly_flat_terrain.xml').read_text(), assets=get_assets())
data  = mujoco.MjData(model)

traj = np.load('npz/g1/dance1_subject1.npz', allow_pickle=True)        # 必须包含 qpos, qvel 两 key
qpos_traj = traj['qpos']          # shape (T, nq)
qvel_traj = traj['qvel']          # shape (T, nv)
qpos_list = []
qvel_list = []
xpos_list = []
geom_xpos_list = []
xquat_list = []
xmat_list = []
cvel_list = []
with mujoco.viewer.launch_passive(model, data) as viewer:
    for t in range(len(qpos_traj)):
        data.qpos[:] = np.concatenate([qpos_traj[t][0:20], [0.0, 0.0], qpos_traj[t][20:25], [0.0, 0.0], qpos_traj[t][25:], [0.0, 0.0]])
        data.qvel[:] = np.concatenate([qvel_traj[t][0:19], [0.0, 0.0], qvel_traj[t][19:24], [0.0, 0.0], qvel_traj[t][24:], [0.0, 0.0]])
        mujoco.mj_forward(model, data)   # 更新所有派生量
        qpos_list.append(data.qpos.copy())
        qvel_list.append(data.qvel.copy())
        xpos_list.append(data.xpos.copy())
        geom_xpos_list.append(data.geom_xpos.copy())
        xquat_list.append(data.xquat.copy())
        xmat_list.append(data.xmat.copy())
        cvel_list.append(data.cvel.copy())
        # print(data.qpos.shape)
        viewer.sync()
        # 控制回放速度，例如 1× 真实时间
        time.sleep(model.opt.timestep)

save_dict = dict(traj)           # 先把原来的 key 都搬过来
save_dict['qpos'] = np.array(qpos_list)
save_dict['qvel'] = np.array(qvel_list)
save_dict['xpos'] = np.array(xpos_list)
save_dict['geom_xpos'] = np.array(geom_xpos_list)
save_dict['xquat'] = np.array(xquat_list)
save_dict['xmat'] = np.array(xmat_list)
save_dict['cvel'] = np.array(cvel_list)
np.savez('npz/g1/dance1_subject1.xpos.npz', **save_dict)  # 覆盖原文件